#include "messiah/dxtc/dxtc_loader.h"
#include "messiah/mcommon.h"
#include "fileio/fileio.h"
#include "fileio/ifile.h"
#include "core/cbuffer.h"

#include "external/DirectXTex/DDS.h"
#include "external/minilzo-2.09/minilzo.h"

#include <iostream>

namespace renderit
{
	namespace messiah
	{
		static constexpr char LzoHeader[] = { 'Z', 'Z', 'Z', 'Z' };

		dxtc_loader::dxtc_loader()
			: scratch_image_(nullptr)
		{

		}

		dxtc_loader::~dxtc_loader()
		{
			if (scratch_image_) {
				delete scratch_image_;
			}
		}

		static bool compress_lzo(core::cbuffer& out, const char * data, int size)
		{
			lzo_uint isize = size;
			lzo_uint osize = isize + isize / 16 + 64 + 3;
			out.reserve(osize + 8);
			uint8 wrkmen[LZO1X_MEM_COMPRESS];
			int ret = lzo1x_1_compress((const unsigned char*)data, isize, (unsigned char*)out.buf() + 8, &osize, wrkmen);
			if (ret != LZO_E_OK) {
				std::cout << "compress_lzo Failed!" << std::endl;
				return false;
			}
			uint8* ptr = (uint8*)out.buf();
			*(uint32*)ptr = *(uint32*)LzoHeader;
			*(uint32*)(ptr + 4) = (uint32)isize;
			out.resize(osize + 8);
			return true;
		}

		static bool save_resource(core::cbuffer& out, const DirectX::Image* images, size_t nimages, const DirectX::TexMetadata& info, const Texture2DInfo& texinfo, const TextureCompressOption& option)
		{
			Texture2DInfo tinfo(texinfo);
			tinfo.MipLevel = info.mipLevels;
			tinfo.Format = EPixelFormat_BC7;
			out.add((char*)&tinfo, sizeof(tinfo));
			int cur = out.tell();

			// Save all slices into a file
			RsTextureInfo texture;
			memset(&texture, 0, sizeof(RsTextureInfo));
			texture.Size = 0;
			texture.Type = ETextureType_2D;
			texture.Format = tinfo.Format;
			texture.MipLevel = info.mipLevels;
			texture.Flags = 0;
			out.add((char*)&texture, sizeof(texture));

			size_t size = sizeof(RsTextureInfo);
			static const char NoneLzo[] = { 'N', 'N', 'N', 'N' };
			size_t extra = sizeof(NoneLzo);
			// Write images
			if (info.dimension != DirectX::DDS_DIMENSION_TEXTURE2D) {
				std::cout << "Unsupported Dimension:" << info.dimension << std::endl;
				return false;
			}

			size_t index = 0;
			for (size_t item = 0; item < info.arraySize; ++item) {
				for (size_t level = 0; level < info.mipLevels; ++level) {
					// for small MIP to larger MIP
					index = info.mipLevels - 1 - level;
					if (index >= nimages) {
						std::cout << "Save Failed, index too big: index:" << index << ", nimages:" << nimages << std::endl;
						return false;
					}
					if (!images[index].pixels) {
						std::cout << "Save Failed, index has no pixels:" << index << std::endl;
						return false;
					}
					
					size_t ddsRowPitch, ddsSlicePitch;
					DirectX::ComputePitch(info.format, images[index].width, images[index].height, ddsRowPitch, ddsSlicePitch, DirectX::CP_FLAGS_NONE);
					if (images[index].slicePitch != ddsSlicePitch) {
						std::cout << "Save Failed, slicePitch not match:" << images[index].slicePitch << " != " << ddsSlicePitch << std::endl;
						return false;
					}

					RsTextureSliceInfo slice;
					memset(&slice, 0, sizeof(slice));
					slice.Size = sizeof(RsTextureSliceInfo) + ddsSlicePitch;
					slice.Width = images[index].width;
					slice.Height = images[index].height;
					slice.Depth = 0;
					slice.PitchInByte = images[index].rowPitch;
					slice.SliceInByte = images[index].slicePitch;

					core::cbuffer buff;
					if (!compress_lzo(buff, (char*)images[index].pixels, ddsSlicePitch)) {
						std::cout << "Save Failed, miplevel:" << index << "/" << info.mipLevels << std::endl;
						return false;
					}
					if (buff.size() < ddsSlicePitch) {
						slice.Size = sizeof(RsTextureSliceInfo) + buff.size();
						out.add((char*)&slice, sizeof(RsTextureSliceInfo));
						out.add(buff.buf(), buff.size());
						std::cout << "Use LZO Compress For MipLevel:" << index << "/" << info.mipLevels << std::endl;
					}
					else
					{
						slice.Size += extra;
						out.add((char*)&slice, sizeof(RsTextureSliceInfo));
						out.add((char*)NoneLzo, extra);
						out.add((char*)images[index].pixels, ddsSlicePitch);
						std::cout << "`NOT` Use LZO Compress For MipLevel:" << index << "/" << info.mipLevels << std::endl;
					}

					size += slice.Size;
				}
			}

			// Rewrite texture 2d info
			texture.Size = size;
			out.add_at(cur, (char*)&texture + offsetof(RsTextureInfo, Size), sizeof(texture.Size));
			return true;
		}

		bool dxtc_loader::convert_dds(const std::string& filepath, const std::string& outpath)
		{
			fileio::ifile * file = fileio::open(filepath, fileio::EOpen_Read | fileio::EOpen_Exist);
			if (!file) {
				std::cout << "Failed To Load File:" << filepath.c_str() << std::endl;
				return false;
			}
			int fsize = file->size();
			char * fbuf = new char[fsize];
			if (file->read(fbuf, fsize) != fsize) {
				file->close();
				std::cout << "Failed To Read File:" << filepath.c_str() << std::endl;
				return false;
			}
			file->close();

			scratch_image_ = new DirectX::ScratchImage();
			HRESULT hr = DirectX::LoadFromDDSMemory(fbuf, fsize, DirectX::DDS_FLAGS_NONE, &tex_meta_data_, *scratch_image_);
			if (FAILED(hr)) {
				std::cout << "Failed To LoadFromDDSMemory:" << filepath.c_str() << std::endl;
				std::cout << "Error:" << hr << std::endl;
				return false;
			}

			Texture2DInfo texinfo;
			texinfo.Width = tex_meta_data_.width;
			texinfo.Height = tex_meta_data_.height;

			if (DirectX::IsPlanar(tex_meta_data_.format)) {
				std::cout << "Unsupported Planar DDS:" << filepath.c_str() << std::endl;
				return false;
			}
			
			// Save
			TextureCompressOption option;
			option.Format = EPixelFormat_BC7;

			const DirectX::Image* img = scratch_image_->GetImage(0, 0, 0);
			size_t nimg = scratch_image_->GetImageCount();

			core::cbuffer buffer;
			bool b = save_resource(buffer, img, nimg, tex_meta_data_, texinfo, option);
			if (!b) {
				return false;
			}

			fileio::ifile* outfile = fileio::open(outpath, fileio::EOpen_Write);
			outfile->write(buffer.buf(), buffer.size());
			return true;
		}
	}
}